# 导入函数库
from jqdata import *
from datetime import date, datetime, timedelta


# 初始化函数，设定基准等等
def initialize(context):
    # 设定沪深300作为基准
    set_benchmark('000300.XSHG')
    # 开启动态复权模式(真实价格)
    set_option('use_real_price', True)
    # 输出内容到日志 log.info()
    log.info('初始函数开始运行且全局只运行一次')
    # 过滤掉order系列API产生的比error级别低的log
    log.set_level('order', 'error')

    # 强制撮合，仅支持限价单。先强制撮合看看效果
    # 强制撮合效果可以啊，但是像一些涨停可能买不进去，误差可能较大
    # set_option("match_by_signal", True)

    # 避免未来数据
    set_option("avoid_future_data", True)

    ### 股票相关设定 ###
    # 股票类每笔交易时的手续费是：买入时佣金万分之三，卖出时佣金万分之三加千分之一印花税, 每笔交易佣金最低扣5块钱
    set_order_cost(OrderCost(close_tax=0.001, open_commission=0.00015, close_commission=0.00015, min_commission=5),
                   type='stock')

    # 将一些需要的参数放到g里面
    # 备选票
    g.option_stocks = []
    g.max_count_list = []
    # 情绪周期，设为3天
    g.emotion_cyc = 3

    # 开盘运行交易
    # run_daily(trade, time='9:30', reference_security='000300.XSHG')
    # 试试33的时候运行，避免一些开盘快速跳水的
    run_daily(trade, time='9:33', reference_security='000300.XSHG')

    run_daily(down_sell, time='10:30', reference_security='000300.XSHG')

    # 尾盘再卖一次
    run_daily(tail_sell, time='14:50', reference_security='000300.XSHG')

    # 收盘后选择股票
    run_daily(get_option_stocks, time='17:00', reference_security='000300.XSHG')


# 交易
def trade(context):
    # 卖出
    sell_trade(context)

    available_cash = context.portfolio.available_cash
    if available_cash < 5000:
        log.info("==========没钱了，今天不动了")
        return

    # 买入
    buy_trade(context)


# 比开盘价低，就卖出
def down_sell(context):
    long_positions_dict = context.portfolio.long_positions
    hold_list = []
    for position in list(long_positions_dict.values()):
        hold_list.append(position.security)

    if len(hold_list) == 0:
        return

    hold_df = get_price(hold_list, end_date=context.current_dt, count=1, frequency='1d',
                        fields=['open'], panel=False, fill_paused=False)
    hold_df.index = hold_df.code

    # 没涨停，卖
    # current_data = get_current_data()
    for position in list(long_positions_dict.values()):
        closeable_amount = position.closeable_amount
        if closeable_amount <= 0:
            continue

        log.info("%s 可卖数量：%d" % (position.security, closeable_amount))
        current_price = position.price
        open_price = hold_df.loc[position.security]['open']
        if current_price < open_price:
            order(position.security, -closeable_amount)
            log.info("=====比开盘低，卖出======stock:%s, current_price:%s, open_price:%s, closeable_amount:%s" % (
            position.security, current_price, open_price, closeable_amount))


# 尾盘卖出
def tail_sell(context):
    long_positions_dict = context.portfolio.long_positions
    hold_list = []
    for position in list(long_positions_dict.values()):
        hold_list.append(position.security)

    if len(hold_list) == 0:
        return

    hold_df = get_price(hold_list, end_date=context.current_dt, count=1, frequency='1d',
                        fields=['high_limit'], panel=False, fill_paused=False)
    hold_df.index = hold_df.code

    '''
    在两个卖出方法的order for循环内部，“position.security”写成了“positon.security”（原来处理hold_list的for “position” 变量名少写了一个i）
    导致order循环中有的地方拿到的股票代码是不对的，但居然在22年收益跑出来还不错（40%左右），这不巧了吗
    然后看交易详情和日志，发现卖出操作不符合策略，百思不得其解，一度以为api哪里用的不对，或者系统有bug，排查了半天。。。。。。
    '''

    # 没涨停，卖
    # current_data = get_current_data()
    for position in list(long_positions_dict.values()):
        closeable_amount = position.closeable_amount
        if closeable_amount <= 0:
            continue

        log.info("%s 可卖数量：%d" % (position.security, closeable_amount))
        # current_data = get_current_data()
        # current_stock = current_data[position.security]
        # current_price = current_stock.last_price
        # high_limit = current_stock.high_limit
        current_price = position.price
        high_limit = hold_df.loc[position.security]['high_limit']
        if current_price < high_limit:
            order(position.security, -closeable_amount)
            log.info("=====尾盘没涨停，卖出======stock:%s, current_price:%s, high_limit:%s, closeable_amount:%s" % (
            position.security, current_price, high_limit, closeable_amount))


# 卖出
def sell_trade(context):
    long_positions_dict = context.portfolio.long_positions
    hold_list = []
    for position in list(long_positions_dict.values()):
        hold_list.append(position.security)

    if len(hold_list) == 0:
        return

    hold_df = get_price(hold_list, end_date=context.current_dt, count=1, frequency='1d',
                        fields=['pre_close'], panel=False, fill_paused=False)
    hold_df.index = hold_df.code

    for position in list(long_positions_dict.values()):
        pre_close = hold_df.loc[position.security]['pre_close']
        # 涨跌幅
        change_percent = round(position.price / pre_close - 1, 2)
        # 绿的，卖了
        # if position.price < pre_close:
        #     order_target(position.security, 0)
        # log.info("======卖出=====stock:%s, pre_close:%s, curr_price:%s, change_percent:%s" % (position.security, pre_close, position.price, change_percent))
        # if (context.current_dt - position.transact_time).days >= 2:
        if change_percent < -0.02:
            order_target(position.security, 0)
            log.info("======跌幅大于2个点，卖出=====stock:%s, pre_close:%s, curr_price:%s, change_percent:%s" % (
            position.security, pre_close, position.price, change_percent))


# 买入
def buy_trade(context):
    stock_size = len(g.option_stocks)
    if stock_size == 0:
        log.info("============无备选股票，今日无买入")
        return
    buy_max_per_stock = round(context.portfolio.total_value * 0.3, 2)
    # 把可用资金平均分给每个备选票
    available_cash = context.portfolio.available_cash
    per_stock_cache = round(available_cash / stock_size, 2)
    buy_cache = min(buy_max_per_stock, per_stock_cache)

    # 盘中获取按天数据时，本以为close就是最新价格，结果居然不能取close，因为close固定为当日收盘价，还没收盘所以没数据
    option_stock_df = get_price(g.option_stocks, end_date=context.current_dt, count=1, frequency='1d',
                                fields=['pre_close', 'open'], panel=False, fill_paused=False)
    option_stock_df.index = option_stock_df.code

    # 连板数量
    limit_count = g.max_count_list[-1]

    # for stock in g.option_stocks:
    #     pre_close = option_stock_df.loc[stock]['pre_close']
    #     open_price = option_stock_df.loc[stock]['open']
    #     # 开盘涨跌幅
    #     change_percent = round(open_price / pre_close - 1, 2)
    #     # 不持有，且红开
    #     # if context.portfolio.positions[stock].closeable_amount <= 0 and open_price >= pre_close:
    #     # if context.portfolio.positions[stock].closeable_amount <= 0 and change_percent > -0.03:
    #     if change_percent > 0:
    #         order_value(stock, buy_cache)
    #         log.info("=====红开，买入====stock:%s, pre_close:%s, open_price:%s, change_percent:%s, limit_count:%s" % (stock, pre_close, open_price, change_percent, limit_count))

    current_data = get_current_data()
    for stock in g.option_stocks:
        current_stock = current_data[stock]
        pre_close = option_stock_df.loc[stock]['pre_close']
        # 通过current_data拿到最新价格
        current_stock = current_data[stock]
        current_price = current_stock.last_price
        # 当前涨跌幅
        change_percent = round(current_price / pre_close - 1, 2)
        if change_percent > 0:
            order_value(stock, buy_cache)
            log.info(
                "=====当前红色，买入====stock:%s, pre_close:%s, current_price:%s, change_percent:%s, limit_count:%s" % (
                stock, pre_close, current_price, change_percent, limit_count))


# 获取备选票
def get_option_stocks(context):
    # 当天涨停的票
    today_limit_stocks = get_today_limit_stocks(context)
    # 当天涨停的票近10天涨停天数
    limit_count_df = get_limit_count(context, list(today_limit_stocks['code']))
    # 取涨停天数最多的票，作为初步备选
    max_limit_count = limit_count_df['count'].max()
    # 把相关数据放到g里
    g.max_count_list.append(max_limit_count)
    # 情绪周期处于上升期，则把股票加入备选, 即：当前最高涨停天数，处于观察期最高
    if len(g.max_count_list) < g.emotion_cyc:
        log.info('=============不够%s天观察期，无备选股票' % (g.emotion_cyc))
        return
    emotion = g.max_count_list[-1] - max(g.max_count_list[-g.emotion_cyc:])
    # 不满足条件，清空g中的备选股票，否则会按前一天的被选票操作买入
    if emotion < 0:
        g.option_stocks = []
        log.info('============观察数据：%s，不符合情绪周期条件' % (str(g.max_count_list[-g.emotion_cyc:])))
    # 满足条件，加入备选
    if emotion == 0:
        limit_count_df = limit_count_df[limit_count_df['count'] == max_limit_count]
        stock_list = list(limit_count_df['code'])
        # 看大佬的策略选择一只市值最小的，试试
        # smallest_stock = get_smallest(context, stock_list)
        # g.option_stocks = [smallest_stock]
        g.option_stocks = stock_list
        log.info(
            '============观察数据：%s，今日备选票：%s' % (str(g.max_count_list[-g.emotion_cyc:]), str(g.option_stocks)))


# 获取当天涨停的票，过滤掉上市不足半年的票
def get_today_limit_stocks(context):
    # 获取所有股票，排除上市不足半年
    all_stocks = get_all_securities(types=['stock'], date=context.current_dt)
    # 半年前的日期
    pre_half_year_date = context.current_dt.date() - timedelta(days=180)
    # 这里的display_name是现在的，回测时根据这个判断是否st不准确
    # all_stocks = all_stocks[(all_stocks['start_date'] < pre_half_year_date) & (~all_stocks['display_name'].str.contains('ST'))]
    all_stocks = all_stocks[all_stocks['start_date'] < pre_half_year_date]

    # 过滤上市不足半年的股票
    today_df = get_price(list(all_stocks.index), end_date=context.current_dt, count=1, frequency='1d',
                         fields=['close', 'high_limit'], panel=False, fill_paused=False)
    today_limit_df = today_df[today_df['close'] == today_df['high_limit']]
    return today_limit_df


# 获取指定股票近10天涨停天数
def get_limit_count(context, stock_list):
    limit_df = get_price(stock_list, end_date=context.current_dt, count=10, frequency='1d',
                         fields=['close', 'high_limit'], panel=False, fill_paused=False)
    limit_df.index = limit_df.code

    count_list = []
    for stock in stock_list:
        df_sub = limit_df.loc[stock]
        limit_days = df_sub[df_sub.close == df_sub.high_limit].close.count()
        count_list.append(limit_days)
    count_df = pd.DataFrame(columns=['code', 'count'])
    count_df['code'] = stock_list
    count_df['count'] = count_list
    # 删除有空值的列
    count_df = count_df.dropna()
    return count_df


# 选择市值最小的票
def get_smallest(context, stock_list):
    q = query(valuation.code, valuation.circulating_market_cap).filter(valuation.code.in_(stock_list)).order_by(
        valuation.circulating_market_cap.asc())
    df = get_fundamentals(q, date=context.current_dt)
    df.index = df.code
    stocks = list(df.code)[0]
    return stocks



